DYMAG: Rethinking Message Passing Using Dynamical-systems-based Waveforms

2023-09-18

AI Summary

Standard graph neural networks (GNNs) rely on local mean aggregation, equivalent to rectangular-wave convolution, limiting their capacity to model complex graph structures. To address this, we propose DYMAG—the first aggregation mechanism that models message-passing waveforms as snapshots of continuous dynamical systems, including the heat equation, wave equation, and Sprott chaotic system—replacing conventional neighborhood averaging with multi-scale, nonlinear waveforms. Theoretically, DYMAG provably identifies connected components, connectivity, and cycle structures without node features and enables cross-scale graph representation learning. Empirically, DYMAG consistently outperforms state-of-the-art GNN baselines across diverse tasks: graph persistence recovery, random graph generation, and prediction of protein, molecular, and materials properties. It significantly enhances structural awareness and generalization performance, demonstrating superior capability in capturing intricate topological patterns.

Abstract

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.

Problem

• Novel message aggregation for graph neural networks
• Dynamic waveforms capture graph structural information
• Outperforms baselines on graph and property prediction tasks

Innovation

• Uses dynamical-systems-based waveforms for message aggregation
• Convolves node features with heat, wave, and chaotic dynamics
• Captures graph structures via multi-scale dynamic snapshots