Hierarchical geometric deep learning enables scalable analysis of molecular dynamics

Molecular dynamics simulations can generate atomically detailed trajectories of complex systems, but analyzing these dynamics can be challenging when systems lack well-established quantitative descriptors (features). Learning features directly from simulation data would obviate manual feature engineering, and graph neural networks are promising architectures for this task. However, the dominant GNN paradigm, message passing, in which information is transferred only between nodes that represent spatially nearby atoms, struggles to capture long-range interactions. Mechanisms that instead allow every node to communicate with every other node, such as attention in transformer architectures, capture long-range interactions but have memory and runtime requirements that scale quadratically with the number of nodes (atoms). Together, these issues limit the use of GNNs for analyzing dynamics of many atoms. Here, we show how a hierarchical scheme can be used to aggregate local information to reduce memory and runtime requirements without sacrificing atomic detail. We demonstrate that this approach opens the door to analyzing simulations of protein-nucleic acid complexes with thousands of residues at full atomic resolution on single GPUs within minutes. For systems with hundreds of residues, for which there are sufficient data to make quantitative comparisons, we show that the approach reduces computational cost while maintaining, and in some cases improving, performance and interpretability.