Self-Conditioned Denoising for Atomistic Representation Learning

Abstract

The success of large-scale pretraining in NLP and computer vision has catalyzed growing efforts to develop analogous foundation models for the physical sciences. However, pretraining strategies using atomistic data remain underexplored. To date, large-scale supervised pretraining on DFT force-energy labels has provided the strongest performance gains to downstream property prediction, out-performing existing methods of self-supervised learning (SSL) which remain limited to ground-state geometries, and/or single domains of atomistic data. We address these shortcomings with Self-Conditioned Denoising (SCD), a backbone-agnostic reconstruction objective that utilizes self-embeddings for conditional denoising across any domain of atomistic data, including small molecules, proteins, periodic materials, and 'non-equilibrium' geometries. When controlled for backbone architecture and pretraining dataset, SCD significantly outperforms previous SSL methods on downstream benchmarks and matches or exceeds the performance of supervised force-energy pretraining. We show that a small, fast GNN pretrained by SCD can achieve competitive or superior performance to larger models pretrained on significantly larger labeled or unlabeled datasets, across tasks in multiple domains. Our code is available at: https://github.com/TyJPerez/SelfConditionedDenoisingAtoms

Figures (4)

• Figure 1: Self-Conditioned Denoising Pretraining
• Figure 2: UMAP of molecule embeddings colored by atom count: SCD results in smoother, less extensive, semantically rich representation. 'Coord' refers to standard node denoising pretraining. See figure \ref{['fig:umap']} for additional plots and details.
• Figure 3: Adaptive layer normalization for SCD: On the left we illustrate a simplified diagram of the original residual layer block in TorchMD-Net (ET). Here, 'X' represents invariant (L0) embeddings while 'V' represents equivariant vector embeddings (L1), and 'C' represents the conditioning embedding. On the right, we show how we modify the original residual block with a conditional scale, shift, and gate of the invariant embeddings. In the gate block the gamma vector is attenuated by a tanh function and multiplied to the output of the ET Attention block. For this work, we use a two layer MLP for conditioning in each layer.
• Figure 4: Pretraining Effects on the Embedding Space. UMAP visualizations of embeddings for 10k QM9 molecules, colored by atom count ($N$). (a) Untrained GNNs exhibit a strong bias towards "extensive" representations, clustering molecules primarily by size ($N$). (b) Standard node denoising smooths the latent manifold but embeddings remain highly extensive. (c--d) SCD pretraining yields a smoother, less extensive (more semantic) embedding space. We observe that downstream performance improves as representations becomes smoother and less extensive (Performance: $a < b < c < d$). 'Molecule embeddings' are created by passing sum pooled atom embeddings through a two-layer MLP 'embedding head'. In the case of 'ET-Coord' the embedding head is untrained. In SCD 'Molecule embeddings' are used for self-embeddings. In our 'coord' model, the molecule 'embedding head' remains untrained.