Sheaf-MPNN:将层神经网络重构为消息传递的 PyTorch Geometric 库
Sheaf-MPNN: A Library for Sheaf Neural Networks as Message Passing
搞图神经网络的朋友可以看看,这个库把 15 种层神经网络统一成消息传递实现,不用再手动拼全局拉普拉斯矩阵,省内存也好扩展。
Sheaf-MPNN 是一个基于 PyTorch Geometric 的开源库,把层神经网络(SNN)的扩散过程重构为局部的逐边传输与聚合,无需显式构建全局稀疏层拉普拉斯矩阵,从而降低内存和计算开销。该库统一实现了 15 个 SNN 家族,支持同质与异质图、有向与无向图,并提供对角、一般、正交、低秩限制映射及基于注意力的预测器。论文证明了这些消息传递形式与原始形式的理论等价性,同时给出覆盖节点级、图级任务和计算效率的统一基准测试套件。
Sheaf-MPNN: A Library for Sheaf Neural Networks as Message Passing
Standard Graph Neural Networks (GNNs) implicitly assume that neighboring nodes share the same feature space, which contributes to oversmoothing and poor performance on heterophilic graphs. Sheaf Neural Networks (SNNs) remove this assumption by giving each node and edge its own vector space and learning restriction maps that transform features as they move between neighbors. These maps define a sheaf Laplacian that drives diffusion over the graph, but existing implementations often assemble it explicitly as a global sparse block matrix, which adds memory and computational overhead and can hinder scalability. In this paper, we introduce Sheaf-MPNN, a PyTorch Geometric library that reformulates sheaf diffusion as local incidence-wise transport and aggregation, eliminating the need to construct the global Laplacian. Sheaf-MPNN implements fifteen families of SNNs under a unified message-passing framework, each supporting both homogeneous and heterogeneous graphs, and directed and undirected settings, with diagonal, general, orthogonal, and low-rank restriction maps and attention-based predictors. We further provide a unified benchmarking suite covering node- and graph-level tasks, heterogeneous graph learning, and computational efficiency, enabling systematic and reproducible comparisons across architectures. We establish theoretical equivalence between these message-passing formulations and their original counterparts and identify the architectural mechanisms needed to preserve their propagation dynamics. This unified perspective provides a principled foundation for implementing, comparing, and extending sheaf-based architectures within standard graph learning frameworks. Source code is available at https://github.com/alessioborgi/sheaf-mpnn.