Scaling Graph Neural Networks to Large Proteins

Abstract

Graph neural network (GNN) architectures have emerged as promising force field models, exhibiting high accuracy in predicting complex energies and forces based on atomic identities and Cartesian coordinates. To expand the applicability of GNNs, and machine learning force fields more broadly, optimizing their computational efficiency is critical, especially for large biomolecular systems in classical molecular dynamics simulations. In this study, we address key challenges in existing GNN benchmarks by introducing a dataset, DISPEF, which comprises large, biologically relevant proteins. DISPEF includes 207,454 proteins with sizes up to 12,499 atoms and features diverse chemical environments, spanning folded and disordered regions. The implicit solvation free energies, used as training targets, represent a particularly challenging case due to their many-body nature, providing a stringent test for evaluating the expressiveness of machine learning models. We benchmark the performance of seven GNNs on DISPEF, emphasizing the importance of directly accounting for long-range interactions to enhance model transferability. Additionally, we present a novel multiscale architecture, termed Schake, which delivers transferable and computationally efficient energy and force predictions for large proteins. Our findings offer valuable insights and tools for advancing GNNs in protein modeling applications.

References in corpus (39)

• Decoupled Weight Decay Regularization
• SchNet - a deep learning architecture for molecules and materials
• ANI-1: An extensible neural network potential with DFT accuracy at force field computational cost
• E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials
• Machine Learning of Accurate Energy-Conserving Molecular Force Fields
• PhysNet: A Neural Network for Predicting Energies, Forces, Dipole Moments and Partial Charges
• Less is more: sampling chemical space with active learning
• Towards Exact Molecular Dynamics Simulations with Machine-Learned Force Fields
• Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds
• Directional Message Passing for Molecular Graphs
• MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields
• Interactive Supercomputing on 40,000 Cores for Machine Learning and Data Analysis
• Equivariant message passing for the prediction of tensorial properties and molecular spectra
• SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
• Coarse Graining Molecular Dynamics with Graph Neural Networks
• Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations
• Cormorant: Covariant Molecular Neural Networks
• GemNet: Universal Directional Graph Neural Networks for Molecules
• Machine Learning Coarse-Grained Potentials of Protein Thermodynamics
• The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials
• End-to-End Differentiable Molecular Mechanics Force Field Construction
• Machine Learning Implicit Solvation for Molecular Dynamics
• Flow-matching -- efficient coarse-graining of molecular dynamics without forces
• TorchMD-NET: Equivariant Transformers for Neural Network based Molecular Potentials
• TorchMD-Net 2.0: Fast Neural Network Potentials for Molecular Simulations
• Adversarial-Residual-Coarse-Graining: Applying machine learning theory to systematic molecular coarse-graining
• MoleculeNet: A Benchmark for Molecular Machine Learning
• Machine-learned molecular mechanics force field for the simulation of protein-ligand systems and beyond
• Contrastive Learning of Coarse-Grained Force Fields
• EspalomaCharge: Machine learning-enabled ultra-fast partial charge assignment
• Equivariant Graph Mechanics Networks with Constraints
• Implicit Solvent Approach Based on Generalised Born and Transferable Graph Neural Networks for Molecular Dynamics Simulations
• Spatial Attention Kinetic Networks with E(n)-Equivariance
• A Survey of Geometric Graph Neural Networks: Data Structures, Models and Applications
• Accurate Molecular Dynamics Enabled by Efficient Physically-Constrained Machine Learning Approaches
• Peering inside the black box: Learning the relevance of many-body functions in Neural Network potentials
• Thermodynamically Informed Multimodal Learning of High-Dimensional Free Energy Models in Molecular Coarse Graining
• Learning data efficient coarse-grained molecular dynamics from forces and noise
• Density-Based Long-Range Electrostatic Descriptors for Machine Learning Force Fields

Cited by in corpus (1)

• Knowledge Distillation of a Protein Language Model Yields a Foundational Implicit Solvent Model