mpnn model
Efficient Parallelization of Message Passing Neural Network Potentials for Large-scale Molecular Dynamics
Machine learning potentials have achieved great success in accelerating atomistic simulations. Many of them relying on atom-centered local descriptors are natural for parallelization. More recent message passing neural network (MPNN) models have demonstrated their superior accuracy and become increasingly popular. However, efficiently parallelizing MPNN models across multiple nodes remains challenging, limiting their practical applications in large-scale simulations. Here, we propose an efficient parallel algorithm for MPNN models, in which additional data communication is minimized among local atoms only in each MP layer without redundant computation, thus scaling linearly with the layer number. Integrated with our recursively embedded atom neural network model, this algorithm demonstrates excellent strong scaling and weak scaling behaviors in several benchmark systems. This approach enables massive molecular dynamics simulations on MPNN models as fast as on strictly local models for over 100 million atoms, vastly extending the applicability of the MPNN potential to an unprecedented scale. This general parallelization framework can empower various MPNN models to efficiently simulate very large and complex systems.
Active Learning Enables Extrapolation in Molecular Generative Models
Antoniuk, Evan R., Li, Peggy, Keilbart, Nathan, Weitzner, Stephen, Kailkhura, Bhavya, Hiszpanski, Anna M.
Although generaIve models hold promise for discovering molecules with opImized desired properIes, they oNen fail to suggest synthesizable molecules that improve upon the known molecules seen in training. We find that a key limitaIon is not in the molecule generaIon process itself, but in the poor generalizaIon capabiliIes of molecular property predictors. We tackle this challenge by creaIng an acIve-learning, closed-loop molecule generaIon pipeline, whereby molecular generaIve models are iteraIvely refined on feedback from quantum chemical simulaIons to improve generalizaIon to new chemical space. Compared against other generaIve model approaches, only our acIve learning approach generates molecules with properIes that extrapolate beyond the training data (reaching up to 0.44 standard deviaIons beyond the training data range) and out-of-distribuIon molecule classificaIon accuracy is improved by 79%. By condiIoning molecular generaIon on thermodynamic stability data from the acIve-learning loop, the proporIon of stable molecules generated is 3.5x higher than the next-best model. More recently, generaIve or inverse-design models have been proposed as a new paradigm for materials discovery due to their ability to efficiently navigate chemical space beyond what is present in exisIng databases. The goal of property-constrained molecular generaIon is to generate novel molecules that possess desirable properIes for the applicaIon of interest. Typically, a ground-truth oracle funcIon is defined for each molecule design task to quanItaIvely assess how well the generated molecules meet the desired molecular properIes. As a means to quickly approximate this oracle funcIon, property predicIon models are used as a surrogate model. These property predicIon models are first trained on a pre-exisIng dataset of molecular properIes to learn the mapping between the chemical structure of the molecules and their target molecular properIes.
Generalization, Expressivity, and Universality of Graph Neural Networks on Attributed Graphs
Rauchwerger, Levi, Jegelka, Stefanie, Levie, Ron
We analyze the universality and generalization of graph neural networks (GNNs) on attributed graphs, i.e., with node attributes. To this end, we propose pseudometrics over the space of all attributed graphs that describe the fine-grained expressivity of GNNs. Namely, GNNs are both Lipschitz continuous with respect to our pseudometrics and can separate attributed graphs that are distant in the metric. Moreover, we prove that the space of all attributed graphs is relatively compact with respect to our metrics. Based on these properties, we prove a universal approximation theorem for GNNs and generalization bounds for GNNs on any data distribution of attributed graphs. The proposed metrics compute the similarity between the structures of attributed graphs via a hierarchical optimal transport between computation trees. Our work extends and unites previous approaches which either derived theory only for graphs with no attributes, derived compact metrics under which GNNs are continuous but without separation power, or derived metrics under which GNNs are continuous and separate points but the space of graphs is not relatively compact, which prevents universal approximation and generalization analysis.
Fine-grained Expressivity of Graph Neural Networks
Bรถker, Jan, Levie, Ron, Huang, Ningyuan, Villar, Soledad, Morris, Christopher
Numerous recent works have analyzed the expressive power of message-passing graph neural networks (MPNNs), primarily utilizing combinatorial techniques such as the $1$-dimensional Weisfeiler-Leman test ($1$-WL) for the graph isomorphism problem. However, the graph isomorphism objective is inherently binary, not giving insights into the degree of similarity between two given graphs. This work resolves this issue by considering continuous extensions of both $1$-WL and MPNNs to graphons. Concretely, we show that the continuous variant of $1$-WL delivers an accurate topological characterization of the expressive power of MPNNs on graphons, revealing which graphs these networks can distinguish and the level of difficulty in separating them. We identify the finest topology where MPNNs separate points and prove a universal approximation theorem. Consequently, we provide a theoretical framework for graph and graphon similarity combining various topological variants of classical characterizations of the $1$-WL. In particular, we characterize the expressive power of MPNNs in terms of the tree distance, which is a graph distance based on the concept of fractional isomorphisms, and substructure counts via tree homomorphisms, showing that these concepts have the same expressive power as the $1$-WL and MPNNs on graphons. Empirically, we validate our theoretical findings by showing that randomly initialized MPNNs, without training, exhibit competitive performance compared to their trained counterparts. Moreover, we evaluate different MPNN architectures based on their ability to preserve graph distances, highlighting the significance of our continuous $1$-WL test in understanding MPNNs' expressivity.
Multi-view self-supervised learning for multivariate variable-channel time series
Brรผsch, Thea, Schmidt, Mikkel N., Alstrรธm, Tommy S.
Labeling of multivariate biomedical time series data is a laborious and expensive process. Self-supervised contrastive learning alleviates the need for large, labeled datasets through pretraining on unlabeled data. However, for multivariate time series data, the set of input channels often varies between applications, and most existing work does not allow for transfer between datasets with different sets of input channels. We propose learning one encoder to operate on all input channels individually. We then use a message passing neural network to extract a single representation across channels. We demonstrate the potential of this method by pretraining our model on a dataset with six EEG channels and then fine-tuning it on a dataset with two different EEG channels. We compare models with and without the message passing neural network across different contrastive loss functions. We show that our method, combined with the TS2Vec loss, outperforms all other methods in most settings.
Neural Message Passing for Objective-Based Uncertainty Quantification and Optimal Experimental Design
Chen, Qihua, Chen, Xuejin, Woo, Hyun-Myung, Yoon, Byung-Jun
Various real-world scientific applications involve the mathematical modeling of complex uncertain systems with numerous unknown parameters. Accurate parameter estimation is often practically infeasible in such systems, as the available training data may be insufficient and the cost of acquiring additional data may be high. In such cases, based on a Bayesian paradigm, we can design robust operators retaining the best overall performance across all possible models and design optimal experiments that can effectively reduce uncertainty to enhance the performance of such operators maximally. While objective-based uncertainty quantification (objective-UQ) based on MOCU (mean objective cost of uncertainty) provides an effective means for quantifying uncertainty in complex systems, the high computational cost of estimating MOCU has been a challenge in applying it to real-world scientific/engineering problems. In this work, we propose a novel scheme to reduce the computational cost for objective-UQ via MOCU based on a data-driven approach. We adopt a neural message-passing model for surrogate modeling, incorporating a novel axiomatic constraint loss that penalizes an increase in the estimated system uncertainty. As an illustrative example, we consider the optimal experimental design (OED) problem for uncertain Kuramoto models, where the goal is to predict the experiments that can most effectively enhance robust synchronization performance through uncertainty reduction. We show that our proposed approach can accelerate MOCU-based OED by four to five orders of magnitude, without any visible performance loss compared to the state-of-the-art. The proposed approach applies to general OED tasks, beyond the Kuramoto model.
Predicting the clinical citation count of biomedical papers using multilayer perceptron neural network
Li, Xin, Tang, Xuli, Cheng, Qikai
The number of clinical citations received from clinical guidelines or clinical trials has been considered as one of the most appropriate indicators for quantifying the clinical impact of biomedical papers. Therefore, the early prediction of the clinical citation count of biomedical papers is critical to scientific activities in biomedicine, such as research evaluation, resource allocation, and clinical translation. In this study, we designed a four-layer multilayer perceptron neural network (MPNN) model to predict the clinical citation count of biomedical papers in the future by using 9,822,620 biomedical papers published from 1985 to 2005. We extracted ninety-one paper features from three dimensions as the input of the model, including twenty-one features in the paper dimension, thirty-five in the reference dimension, and thirty-five in the citing paper dimension. In each dimension, the features can be classified into three categories, i.e., the citation-related features, the clinical translation-related features, and the topic-related features. Besides, in the paper dimension, we also considered the features that have previously been demonstrated to be related to the citation counts of research papers. The results showed that the proposed MPNN model outperformed the other five baseline models, and the features in the reference dimension were the most important.