Sheaf-MPNN: A Library for Sheaf Neural Networks as Message Passing
Authors: Alessio Borgi, Mario Severino, Luke Braithwaite, Steve Azzolin, Rémi Bourgerie, Olga Zaghen, Emanuele Mule, Riccardo Ali, Leo Minh Kustermann, Edoardo Coppola, Antonio Purificato, Stefano Fiorini, Christopher Irwin, Filippo Spiga, Fabrizio Silvestri, Michael Bronstein, Pietro Liò
Organizations: Department of Computer Science and Technology, University of Cambridge, Cambridge, UK · Department of Computer, Control and Management Engineering, Sapienza University of Rome, Rome, Italy · Department of Information Engineering and Computer Science, University of Trento, Trento, Italy · KTH Royal Institute of Technology, Stockholm, Sweden · AMLab, University of Amsterdam, Amsterdam, Netherlands · Télécom Paris, Institute Polythechnique de Paris, Paris, France · ETH Zürich · Independent researcher · NVIDIA · Department of Computer Science, University of Oxford, Oxford, UK · AITHYRA, Vienna, Austria
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.
Figures & tables
Figure 1: Sheaf Message Passing Neural Network (SheafMPNN). Neighbourhood update at node 0 (red), with neighbouring nodes in blue. (Left) A standard MPNN computes messages m0i (green) from neighbouring node states and aggregates them at node 0 . (Right) A SheafMPNN transports both endpoint representations into the shared edge stalk F(e0i) (boxes) through restriction maps (dashed grey arrows). The function ψ computes their interaction, and the resulting message is pulled back through \F0⊴e0i⊤ and aggregated at node 0 (solid black arrows).
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Module
Responsibility
sheaf_mpnn.base_conv
Shared stalk transforms, map-predictor inputs, and propagation interfaces.
sheaf_mpnn.<family>
Architecture-specific convolutions and end-to-end models.
sheaf_mpnn.utils
Normalization, orthogonal parameterizations, connection maps, and training utilities.
sheaf_mpnn.sheaf_pool
Graph-level pooling of stalk-valued representations.
hetsheaf_mpnn
Type-conditioned layers and model wrappers.
exp.config , exp.registries
Typed experiment settings and factories for models, datasets, and presets.
Appendix
Table 1: Sheaf-MPNN Main Modules.
Family
Model type
Map and model options
Neural Sheaf Diffusion
nsd
Six
Polynomial NSD
poly_nsd
Six
Connection-Laplacian NSD
conn_nsd
Precomputed maps
Deep NSD
dnsd
Six
Neural Sheaf Propagation
nsp
Six
Sheaf adjacency
sheaf_adj
Six
Appendix
Table 2: Homogeneous Architecture Families. The table lists the 15 homogeneous families and their corresponding model types and restriction-map options. In the last column, “six” denotes the six NSD restriction-map variants of Table 3 , and “four” the diagonal, general, orthogonal, and low-rank maps.
Variant
Restriction parameterization
Outputs
DIAGONAL
diag(w) , with w∈(−1,1)d
d
GENERAL
Dense matrix with entries in (−1,1)
d2
GENERAL_ATTENTION
I−softmax(P) , with row-wise softmax
d2
ORTHOGONAL
Cayley or Householder construction
(2d) or d2
ORTHOGONAL_ATTENTION
Cayley transform of the skew part of I−softmax(P)
d2
LOW_RANK
AB⊤ , with A,B∈Rd×r
2dr
Appendix
Table 3: NSD Restriction-map Variants . The last column gives the number of predictor outputs per restriction map, excluding fixed channels and edge weights. Low-rank maps need fewer outputs, but their products Sue and Cuve are still stored as dense d×d blocks.
Option
Effect
variant , rank , orth_strategy
Select the restriction-map variant and its parameterization.
alpha , learn_alpha , add_eps
Set the diffusion step and enable the learnable skip scaling.
use_edge_weights
Enable learned symmetric edge weights for orthogonal maps.
add_lp , add_hp
Append a fixed channel with signless ( D+A ) or standard ( D−A ) graph-Laplacian blocks, respectively.
sparse_learner
Sum the paired contexts over the stalk axis, reducing the predictor width from 2dtotf to 2f for flattened stalk inputs.
second_linear
Add a second encoder layer.
Appendix
Table 4: Options of NSDModel beyond the core dimensions in_channels , out_channels , stalk_dim , hidden_dim , and num_layers .
Sheaf Neural Networks (SNNs) were introduced as an extension of Graph Convolutional Networks to address oversmoothing on heterophilous graphs by attaching a sheaf to the input graph and replacing the adjacency-based operator with a sheaf Laplacian defined by (learnable) restriction maps. Prior work motivates this design through theoretical properties of sheaf diffusion and the kernel of the sheaf Laplacian, suggesting that suitable non-identity restriction maps can avoid representations converging to constants across connected components. Since oversmoothing can also be mitigated through residual connections and normalization, we revisit a trivial sheaf construction to ask whether the additional complexity of learning restriction maps is necessary. We introduce an Identity Sheaf Network baseline, where all restriction maps are fixed to the identity, and use it to ablate the empirical improvements reported by sheaf-learning architectures. Across five popular heterophilic benchmarks, the identity baseline achieves comparable performance to a range of SNN variants. Finally, we introduce the Rayleigh quotient as a normalized measure for comparing oversmoothing across models and show that, in trained networks, the behavior predicted by the diffusion-based analysis of SNNs is not reflected empirically. In particular, Identity Sheaf Networks do not appear to suffer more significant oversmoothing than their SNN counterparts.
Ferran Hernandez Caralt, Mar Gonzàlez i Català, Adrián Bazaga +1
Department of Computer Science and Technology, University of Cambridge, Cambridge, United Kingdom · Microsoft
Sheaf Neural Networks (SNNs) generalize Graph Neural Networks (GNNs) by replacing scalar node signals with stalk-valued signals and by using restriction maps to measure compatibility across edges. Unlike standard graph diffusion, which encourages neighboring node features to become similar, sheaf diffusion promotes consistency through the restriction maps and can therefore model more general relationships between neighboring nodes. However, existing sheaf neural architectures mainly operate at a fixed graph resolution and do not provide a principled pooling mechanism for building hierarchical representations. In this paper, we introduce Hierarchical Sheaf Pool (HiSP), a sheaf-aware pooling framework based on local spectral coarsening. Given a partition of the graph, HiSP constructs each coarse stalk by projecting fine stalk-valued features onto the low-frequency eigenmodes of the cluster-internal sheaf Laplacian. These local modes define a cochain-level prolongation map, which allows the fine sheaf energy to be represented on the coarse space through a Galerkin operator. We further analyze the approximation induced by coarsening by separating truncation loss, due to discarded local modes, from realization loss, due to representing the projected operator as a coarse sheaf. Finally, we implement HiSP as a GNN pooling layer compatible with SNNs and provide a PyG implementation supporting batching, lifted sheaf Laplacians, and hierarchical architectures.
Dionisia Naddeo, Carlo Abate, Pietro Liò +2
UiT The Arctic University of Norway · UiT The Arctic University of Norway NORCE Norwegian Research Centre AS
Graph Neural Networks (GNNs) have become the de facto standard for learning on relational data. While traditional GNNs' message passing is well suited for vector-valued node features, there are cases in which node features are better represented by probability distributions than real vectors. Concretely, when node features are Gaussians, characterized by a mean and a covariance matrix, naively concatenating their parameters into a single vector and applying standard message passing discards the geometric and algebraic structure that governs means and covariances. We propose Gaussian Sheaf Neural Networks (GSNNs), a principled framework that incorporates these inductive biases into graph-based learning. Building on the theory of cellular sheaves, we derive a new Laplacian operator that generalizes the sheaf Laplacian to this setting and preserves its key properties. We complement our theoretical contributions with experiments on synthetic and real-world data that illustrate the practical relevance of GSNNs.
André Ribeiro, Ana Luiza Tenório, Tiago da Silva +1