PaperlayerView exact revision on arXiv

Source-linked AI summary

Pre-training via Denoising for Molecular Property Prediction

Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzalez, Peter Battaglia, Razvan Pascanu, Jonathan Godwin

arXiv:2206.00133v2cs.LGq-bio.BMstat.ML

TL;DR

Molecular property prediction from 3D structures often has limited labeled data, leaving effective pre-training for this setting unresolved. The paper denoises equilibrium structures from large datasets and interprets the objective as learning a force field. This substantially improves performance across multiple benchmarks, including a new state-of-the-art on most QM9 targets.

  • Problem

    Effective pre-training for molecular property prediction from 3D structures remains unresolved despite limited data and the availability of large 3D structure datasets.

  • Method

    The paper pre-trains molecular GNNs by predicting Gaussian noise added to equilibrium structures, interpreting denoising as force-field learning under a Boltzmann-distribution approximation.

  • Results

    Pre-training via denoising significantly improves performance on multiple challenging datasets and achieves a new state-of-the-art on 10 out of 12 QM9 targets.

  • Takeaways & Limitations

    Denoising enables large datasets of 3D molecular structures to provide transferable representations for downstream molecular property prediction.

  • Takeaways & Limitations

    Transfer depends on the relationship between upstream and downstream datasets, and molecular networks lack a settled standard architecture that would make pre-trained weights broadly reusable.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper addresses the open problem of pre-training representations for molecular property prediction from 3D structures. It introduces denoising pre-training, interprets it as force-field learning, and reports state-of-the-art results across molecular benchmarks.

  • Motivation and approach: Large 3D molecular datasets can be used for self-supervised denoising pre-training to improve downstream molecular property prediction.The approach targets downstream tasks that also use 3D structures, including settings where downstream datasets are much smaller than upstream datasets.
  • Motivation and approach: The denoising objective is equivalent to learning a particular molecular force field through the connection between denoising and score-matching.This provides a force-field interpretation of how denoising can aid representation learning.
  • Main findings: Pre-training via denoising significantly improves performance on multiple challenging datasets and reaches a new state-of-the-art on 10 out of 12 QM9 targets.The evaluated datasets vary in size, task type, and molecular composition.
  • Main findings: The study also contributes GNS-TAT, applying Tailored Activation Transformation to Graph Network Simulators as a complementary improvement to pre-training.Figure 1 reports GNS-TAT pre-trained via denoising on PCQM4Mv2 outperforming prior work on QM9.

2 RELATED WORK

Related work spans graph pre-training, denoising-based representation learning, score-matching, and equivariant architectures for 3D molecular prediction. The paper distinguishes its approach by applying denoising pre-training directly to 3D molecular point clouds.

  • Pre-training and 3D architectures: Existing molecular pre-training methods largely operate on 2D graphs, whereas this paper targets 3D point clouds of atomic nuclei.The downstream setting uses three-dimensional molecular structures as inputs.
  • Pre-training and 3D architectures: Contrastive and bootstrapping methods learn graph representations from contrasting views or related objectives.These approaches form one branch of recent graph-based pre-training research.
  • Denoising and score-matching: Denoising autoencoders learn representations by mapping corrupted inputs back to their originals, and denoising has also improved GNN and molecular prediction performance.The paper builds on denoising as both a representation-learning mechanism and an auxiliary task.
  • Denoising and score-matching: Denoising is closely connected to score-matching, and the paper uses this connection to interpret structure denoising as learning a molecular force field.Score-matching has also been used in generative modeling.
  • Pre-training and 3D architectures: Equivariant neural networks improve 3D molecular prediction by incorporating roto-translational inductive biases so outputs remain invariant to translations and rotations.Examples include invariant geometric features such as inter-atomic distances and angles.

3 METHODOLOGY

The method pre-trains a GNN by predicting Gaussian noise added to equilibrium molecular coordinates, framing denoising as force-field learning under a smoothed Boltzmann-distribution approximation. It also extends GNS with TAT and uses denoising-related auxiliary objectives for representation learning.

  • 3.1 PROBLEM SETUP: Molecular property prediction maps atomic numbers and 3D positions to scalar labels using a backbone representation and vertex-level or graph-level decoder.The input is a set of atoms with atomic number a_i and position p_i ∈ R3, alongside a scalar target y.
  • 3.2 PRE-TRAINING VIA DENOISING: Pre-training samples equilibrium structures and trains a GNN to return per-vertex predictions for denoising.The upstream dataset consists of molecular structures, and the model parameters are optimized through the denoising objective.
  • 3.2 PRE-TRAINING VIA DENOISING: Gaussian noise is independently added to each atomic position as ˜p_i = p_i + σϵ_i, with ϵ_i ∼ N(0, I3).The noise scale σ is a tunable hyperparameter, and the model predicts the noise rather than the original coordinates.
  • 3.2.1 DENOISING AS LEARNING A FORCE FIELD: The force-field motivation follows from treating structures as Boltzmann-distributed and approximating the unknown physical distribution with Gaussians centered on equilibrium structures.The equilibrium structures are available, while the energy function and physical distribution are unknown or expensive to evaluate directly.
  • 3.2.1 DENOISING AS LEARNING A FORCE FIELD: The resulting force field is computationally convenient and yields a score-matching objective equivalent to the denoising loss up to a rescaling absorbed by the network.The Gaussian mixture smooths the distribution around the observed equilibrium structures.
  • 3.2.2 NOISY NODES: DENOISING AS AN AUXILIARY LOSS: Noisy Nodes adds denoising as an auxiliary loss on the downstream dataset, while this paper applies denoising pre-training on a separate large unlabeled structure dataset.The separate upstream objective is intended to learn higher-quality representations before downstream fine-tuning.
  • 3.3 GNS AND GNS-TAT: GNS-TAT combines Graph Net Simulator with Tailored Activation Transformation using Edge-Delta initialization to make the architecture compatible with TAT.The paper presents this architectural improvement as complementary to pre-training.

4 EXPERIMENTS

The experiments test denoising pre-training across molecular datasets that differ in size, composition, and labeling methodology, with transfer depending on upstream–downstream compatibility. On QM9, pre-training substantially improves performance, while transfer to OC20 is limited by dataset differences.

  • Experimental design: The evaluation spans four molecular datasets varying in size, composition, and DFT- versus CCSD(T)-based labeling.The experiments address both whether denoising pre-training helps and how its benefit depends on the relationship between upstream and downstream datasets.
  • Datasets: PCQM4Mv2 provides 3.4 million equilibrium 3D organic molecules for self-supervised pre-training, while QM9, OC20, and DES15K support downstream evaluation.PCQM4Mv2 labels are unused because denoising requires only structures; DES15K contains about 15,000 dimer geometries.
  • Dataset relationships: QM9 is fully covered by PCQM4Mv2, whereas less than 30% of OC20 elements overlap, suggesting weaker transfer to OC20.OC20 contains inorganic surface lattices absent from PCQM4Mv2, while DES15K also has less element overlap than QM9.
  • QM9 results: Pre-training improved all QM9 targets and achieved state-of-the-art performance on 10 of 12 targets.The comparison used GNS-TAT with Noisy Nodes, contrasting random initialization with parameters pre-trained on PCQM4Mv2 and averaging over three seeds.
  • QM9 results: GNS-TAT performed worse than other models on the electronic spatial extent target, possibly because its optimal denoising noise scale differs across targets.The stated explanation is tentative and concerns the fixed hyperparameter choice used across QM9 targets.
  • OC20 results: On OC20 IS2RE, PCQM4Mv2 pre-training provided no validation benefit, whereas OC20 pre-training accelerated convergence without changing final performance.The authors associate the limited cross-dataset transfer with differences in dataset nature and element overlap.

5 ANALYSIS

The analysis examines whether denoising pre-training transfers across architectures, data regimes, model scales, force-prediction tasks, and frozen representations. Benefits persist beyond the original architecture and can be especially valuable with limited downstream data or smaller models.

  • Pre-training a different architecture: Pre-training improves TorchMD-NET performance on QM9 HOMO and LUMO targets, yielding state-of-the-art results despite the architecture differing from GNS-TAT.The authors report that their pre-training results also outperform the published TorchMD-NET results used for comparison.
  • Varying dataset sizes: Pre-training improves HOMO performance for every tested QM9 dataset size, with the test-MAE gap increasing as downstream training data decreases.This analysis directly compares GNS-TAT with Noisy Nodes from scratch against the same model using pre-trained parameters.
  • Varying model size: Pre-trained models continue to benefit from scaling from 10 million to 1.2 billion parameters, and a model smaller than 650 million parameters outperforms a 130-million-parameter model trained from scratch.The comparison uses four pre-trained GNS models and one 130-million-parameter scratch model on OC20 S2EF.
  • Force prediction: Pre-training improves TorchMD-NET force prediction on MD17 aspirin and similarly improves force prediction on OC20.The motivation is that denoising structures corresponds to learning an approximate force field from equilibrium structures.
  • Freezing pre-trained parameters: With the backbone frozen, a decoder trained on pre-trained features reaches 40 meV test MAE versus more than 100 meV with randomly initialized features.Training only the decoder is also substantially faster: one batch takes 15 ms on a single V100 GPU, reported as 50× faster than training the full network.

6 LIMITATIONS & FUTURE WORK

Pre-training improves performance across tasks, but transfer depends on upstream–downstream dataset relationships and incurs computational cost. Molecular models also lack a standard architecture, limiting reuse of shared weights.

  • Training only the decoder performs significantly better with pre-trained features than with random features.Figure 5 compares fixed-backbone pre-trained and randomly initialized representations on the decoder-only evaluation.
  • Pre-trained models can be shared for downstream fine-tuning, but molecular networks vary widely and lack a settled standard architecture.This makes pre-trained weights less reusable than in vision and NLP.
  • Transfer depends on which aspects of the upstream data distribution match the downstream distribution.The paper identifies this relationship as an important direction for future work.
  • Pre-training large models incurs additional computational cost, although 3D molecular prediction does not require the scale of large NLP and vision models.

7 CONCLUSION

The paper presents denoising pre-training for 3D molecular structures as a way to learn useful representations and improve downstream molecular property prediction. It connects denoising to force-field learning and reports improved performance across tasks, including new state-of-the-art results in some cases such as QM9.

  • Denoising 3D molecular structures is used to pre-train neural networks for downstream molecular property prediction.
  • The denoising objective is equivalent to learning a molecular force field, providing an interpretation for the learned representations.
  • Pre-training improves performance on various downstream molecular property prediction tasks and sets a new state of the art in some cases, such as QM9.
  • The base model is a Graph Net Simulator with an encoder, repeated message-passing processor blocks, and a decoder.
  • Graph-level predictions apply an MLP to vertex features, sum the resulting values over vertices, and process the aggregate with another MLP.

A.2 GNS WITH TAILORED ACTIVATION TRANSFORMATION (GNS-TAT)

GNS-TAT adapts Tailored Activation Transformation to graph networks whose structure violates TAT’s assumptions. The method uses Edge-Delta initialization to enable the transformation while preserving edge information, with related architectural modifications.

  • Motivation: Oversmoothing causes vertex and edge features to converge with depth and is associated with training difficulties in GNNs.
  • Motivation: GNS violates TAT assumptions because vertex updates sum edge features and the network must preserve multiple inputs independently.
  • Edge-Delta initialization: Edge-Delta initializes weights multiplying incoming vertex features in edge updates to zero, making those connections effectively absent for variance calculations.
  • Edge-Delta initialization: With Edge-Delta, edge features propagate independently through an effective deep residual edge network.
  • TAT modifications: GNS-TAT replaces activations with Tailored ReLUs and uses weighted shortcut and non-shortcut branches with weights 0.9 and 1 − 0.9^2, respectively.
  • Initialization comparison: Edge-Delta performed best among the tested Edge-Delta and Vertex-Delta initialization variants, possibly because molecular tasks rely more on edge-encoded distances.

B DENOISING AS LEARNING A FORCE FIELD

The paper formulates denoising equilibrium structures as force-field learning through score matching. Because translation invariance makes the physical density improper, the analysis restricts structures to a mean-centered subspace before using Gaussian perturbations.

  • The molecular force field is the mapping x 7→−∇xE(x), which equals the score ∇x log pphysical(x) under the Boltzmann distribution.Here pphysical(x) ∝ exp(−E(x)).
  • The physical density is improper because translation invariance makes E(x + t) = E(x), so its normalizing integral diverges.
  • Restricting structures to the mean-centered subspace V makes the physical distribution normalizable.V has dimension 3N − 3 and contains structures whose mean position is zero.
  • The analysis uses mean-centered equilibrium structures and Gaussian distributions qσ(˜x | xi) defined on V.
  • The denoising objective uses the conditional score ∇˜x log qσ(˜x | x) = (x − ˜x)/σ^2, and its noise should be mean-centered.
  • Pre-training can support applications in materials design, chemistry, and drug discovery, while computational cost motivates sharing embeddings and using smaller diverse upstream subsets.

D DATASETS

The study uses several molecular datasets spanning equilibrium structures, property labels, molecular compositions, dataset sizes, and labeling methods to examine denoising pre-training across downstream settings.

  • Upstream dataset: PCQM4Mv2 contains 3,378,606 equilibrium 3D organic molecules with approximately 30 atoms on average and 22 unique elements.The structures include atom types and coordinates calculated using DFT.
  • Downstream datasets: QM9 contains approximately 130,000 small organic molecules with 12 molecular-property labels used for fine-tuning.The molecules contain up to nine heavy carbon, nitrogen, oxygen, and fluorine atoms.
  • Downstream datasets: OC20 contains trajectories of interacting surfaces and adsorbates relevant to catalyst discovery, with tasks covering relaxed energies, relaxed structures, and energy-force prediction.The experiments use its IS2RE setting for predicting relaxed energy from an initial structure.
  • Downstream datasets: DES15K contains around 15,000 interacting molecule pairs labeled with CCSD(T) interaction energies.Its labels use a higher-level electronic-structure method than DFT.
  • Dataset relationship: Denoising equilibrium structures provides a label-agnostic pre-training model that can be fine-tuned across targets and can transfer from DFT-generated structures to CCSD(T)-labeled downstream tasks.The authors emphasize reuse across differing downstream needs and a demonstrated benefit on DES15K.

E EXPERIMENT SETUP AND COMPUTE RESOURCES

The experiments combine denoising pre-training and fine-tuning across several architectures and datasets, with reported compute spanning TPU and GPU resources and requiring hyperparameter sweeps for noise and loss settings.

  • Compute resources: GNS-TAT training for QM9, PCQM4Mv2, and DES15K used 16 TPU v3 devices, while evaluation used one V100 device.OC20 GNS training used 8 TPU v4 devices, except the 1.2-billion-parameter model, which used 64 TPU v4 devices.
  • Compute resources: OC20 GNS experiments included a 1.2-billion-parameter variant trained on 64 TPU v4 devices.Other OC20 GNS models were trained on 8 TPU v4 devices.
  • Compute resources: TorchMD-NET pre-training on PCQM4Mv2 used three GPUs and took roughly 24 hours, while QM9 training used two RTX 2080Ti GPUs.Fine-tuning used the same implementation with defaults except for the learning-rate schedule.
  • Hyperparameter optimization: Effective denoising pre-training required sweeping noise values and loss coefficients for denoising and atom-type recovery.For GNS/GNS-TAT, the authors tuned approximately five candidate values for each noise setting on PCQM4Mv2 and QM9.
  • Reported downstream results: TorchMD-NET pre-training further improved DES15K interaction-energy performance beyond its from-scratch performance.The same section reports that TorchMD-NET outperformed GNS-TAT when trained from scratch, before the additional pre-training boost.

G.2 COMPARISON OF OC20 IS2RE PRE-TRAINING PERFORMANCE WITH OTHER ARCHITECTURES

On OC20, the study compares GNS variants with other architectures for IS2RE and evaluates denoising pre-training for force prediction, reporting state-of-the-art performance and broad force-prediction gains.

  • IS2RE comparison: GNS achieves state-of-the-art performance on each of the four OC20 validation sets for direct IS2RE prediction.Table 8 compares GNS variants with other proposed architectures.
  • IS2RE pre-training: The model pre-trained on OC20 achieves state-of-the-art performance with faster convergence than the other GNS variants.At convergence, it performs equally well as the model trained from scratch with noisy nodes.
  • Force prediction: Pre-training improves force prediction over training from scratch for both mean absolute error and cosine similarity across all four OC20 validation datasets.Mean absolute error is lower-is-better, whereas cosine similarity is higher-is-better.
  • Force prediction: Force-prediction improvements from denoising pre-training reach up to 15%.The comparison uses models pre-trained by denoising equilibrium structures in OC20.
Loading 2206.00133v2…