[Paper Review] Equivariant Graph Neural Networks for 3D Macromolecular Structure

Bowen Jing, Stephan Eismann+2 more|arXiv (Cornell University)|Jun 7, 2021

Machine Learning in Materials Science14 references47 citations

TL;DR

This paper extends geometric vector perceptron-based GNNs to atomic-level macromolecular structures with vector gating, achieving strong results on ATOM3D tasks and demonstrating transfer learning benefits.

ABSTRACT

Representing and reasoning about 3D structures of macromolecules is emerging as a distinct challenge in machine learning. Here, we extend recent work on geometric vector perceptrons and apply equivariant graph neural networks to a wide range of tasks from structural biology. Our method outperforms all reference architectures on three out of eight tasks in the ATOM3D benchmark, is tied for first on two others, and is competitive with equivariant networks using higher-order representations and spherical harmonic convolutions. In addition, we demonstrate that transfer learning can further improve performance on certain downstream tasks. Code is available at https://github.com/drorlab/gvp-pytorch.

Motivation & Objective

• Motivate robust learning on 3D macromolecular structures that is invariant to pose and orientation.
• Develop an atomic-level, equivariant GNN using geometric vector perceptrons with vector gating.
• Benchmark the model against CNN, GNN, and ENN baselines on the ATOM3D tasks.
• Explore transfer learning to improve performance on data-poor macromolecular datasets.

Proposed method

• Represent macromolecules as graphs with atoms as nodes and edges for pairs within 4.5 Å.
• Use GVP-GNNs with scalar and vector features to preserve rotation equivariance during message passing.
• Introduce vector gating to propagate information from scalar to vector channels in atomic-level inputs.
• Extend GVP-GNN to atomic-level structures with five layers and mean pooling for downstream tasks.
• Compare against CNN, GNN, and ENN baselines on ATOM3D benchmarks.

Experimental results

Research questions

• RQ1 Can equivariant message-passing with vector features effectively model atomic-level macromolecular structures?
• RQ2 Do vector-gated GVP-GNNs achieve competitive performance across diverse ATOM3D tasks compared to existing baselines?
• RQ3 Is transfer learning from data-rich tasks beneficial for data-poor macromolecular tasks?
• RQ4 How does the atomic-level extension of GVP-GNN compare to higher-order equivariant networks in this domain?

Key findings

• The vector-gated GVP-GNN is the top method on three ATOM3D tasks and tied for first on two others.
• It is competitive with higher-order ENN architectures despite using only first-order vector representations.
• Transfer learning from SMP and RES improves performance on several downstream tasks (PPI, MSP, PSR).
• Pretraining speeds up learning on some tasks, though it may overfit on others (e.g., LBA).
• The approach demonstrates strong performance across diverse macromolecular structure tasks and provides a universal approximation property extension.