DYMAG: Rethinking Message Passing Using Dynamical-systems-based Waveforms
This work addresses the need for more expressive graph neural networks for researchers and practitioners in fields like bioinformatics and materials science, though it is incremental as it builds on existing message-passing frameworks.
The authors tackled the problem of limited expressiveness in standard message-passing neural networks by introducing DYMAG, which uses dynamical-systems-based waveforms for message aggregation, resulting in improved performance on graph recovery, parameter generation, and property prediction tasks across synthetic and real-world benchmarks.
We present DYMAG, a graph neural network based on a novel form of message aggregation. Standard message-passing neural networks, which often aggregate local neighbors via mean-aggregation, can be regarded as convolving with a simple rectangular waveform which is non-zero only on 1-hop neighbors of every vertex. Here, we go beyond such local averaging. We will convolve the node features with more sophisticated waveforms generated using dynamics such as the heat equation, wave equation, and the Sprott model (an example of chaotic dynamics). Furthermore, we use snapshots of these dynamics at different time points to create waveforms at many effective scales. Theoretically, we show that these dynamic waveforms can capture salient information about the graph including connected components, connectivity, and cycle structures even with no features. Empirically, we test DYMAG on both real and synthetic benchmarks to establish that DYMAG outperforms baseline models on recovery of graph persistence, generating parameters of random graphs, as well as property prediction for proteins, molecules and materials. Our code is available at https://github.com/KrishnaswamyLab/DYMAG.