Bundle Neural Networks for message diffusion on graphs

2024-05-24

arXiv.org

AI Summary

Standard message-passing paradigms in Graph Neural Networks (GNNs) suffer from inherent limitations including over-smoothing, over-compression, and restricted node-level expressivity. To address these issues, we propose Bundle Neural Networks (BuNN), the first GNN framework that incorporates flat vector bundles—concepts from differential geometry—into graph representation learning. BuNN establishes a vector-bundle-based message diffusion mechanism, enabling continuous-dynamic modeling via diffusion-type partial differential equations (PDEs). We theoretically prove that BuNN achieves universal node-level expressivity and alleviates over-compression from a geometric perspective. Methodologically, BuNN integrates orthogonal mapping-driven graph augmentation, PDE-inspired feature evolution, a sheaf-compatible discretization scheme, and injective positional encoding. Under both transductive and inductive settings, BuNN attains state-of-the-art performance across multiple standard benchmarks. Synthetic experiments further demonstrate its strong robustness against over-smoothing and over-compression.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Computer Vision: Diffusion Models for Vision
• Reasoning under Uncertainty: Graphical Models

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Social Networks and Social Media: Influence propagation, information diffusion, and the prediction on networks
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

The dominant paradigm for learning on graph-structured data is message passing. Despite being a strong inductive bias, the local message passing mechanism suffers from pathological issues such as over-smoothing, over-squashing, and limited node-level expressivity. To address these limitations we propose Bundle Neural Networks (BuNN), a new type of GNN that operates via message diffusion over flat vector bundles - structures analogous to connections on Riemannian manifolds that augment the graph by assigning to each node a vector space and an orthogonal map. A BuNN layer evolves the features according to a diffusion-type partial differential equation. When discretized, BuNNs are a special case of Sheaf Neural Networks (SNNs), a recently proposed MPNN capable of mitigating over-smoothing. The continuous nature of message diffusion enables BuNNs to operate on larger scales of the graph and, therefore, to mitigate over-squashing. Finally, we prove that BuNN can approximate any feature transformation over nodes on any (potentially infinite) family of graphs given injective positional encodings, resulting in universal node-level expressivity. We support our theory via synthetic experiments and showcase the strong empirical performance of BuNNs over a range of real-world tasks, achieving state-of-the-art results on several standard benchmarks in transductive and inductive settings.

Problem

• Addresses over-smoothing in graph neural networks
• Mitigates over-squashing via message diffusion
• Enhances node-level expressivity universally

Innovation

• Uses Bundle Neural Networks for message diffusion
• Mitigates over-smoothing via diffusion-type PDE
• Achieves universal node-level expressivity