Using pretrained graph neural networks with token mixers as geometric featurizers for conformational dynamics

TL;DR

This paper introduces geom2vec, a method that uses pretrained graph neural networks as universal geometric featurizers to characterize molecular conformational dynamics, reducing manual effort and improving robustness.

Contribution

The paper presents a novel approach combining pretrained GNNs with token mixers as geometric featurizers, enabling transferable and interpretable structural representations for molecular dynamics.

Findings

Pretrained GNNs effectively capture structural features of molecular conformations.

Decoupling GNN training from downstream tasks allows analysis of larger molecular graphs.

Geom2vec improves robustness and reduces manual feature engineering in molecular simulations.

Abstract

Identifying informative low-dimensional features that characterize dynamics in molecular simulations remains a challenge, often requiring extensive manual tuning and system-specific knowledge. Here, we introduce geom2vec, in which pretrained graph neural networks (GNNs) are used as universal geometric featurizers. By pretraining equivariant GNNs on a large dataset of molecular conformations with a self-supervised denoising objective, we obtain transferable structural representations that are useful for learning conformational dynamics without further fine-tuning. We show how the learned GNN representations can capture interpretable relationships between structural units (tokens) by combining them with expressive token mixers. Importantly, decoupling training the GNNs from training for downstream tasks enables analysis of larger molecular graphs (such as small proteins at all-atom resolution) with limited computational resources. In these ways, geom2vec eliminates the need for manual feature selection and increases the robustness of simulation analyses.